import sys
sys.path.insert(0, "/home/ninetripod/catkin_ws/src/uav/scripts")
sys.path.insert(0, "/home/ninetripod/catkin_ws/src/uav/scripts/visual")
sys.path.insert(0, "/home/ninetripod/catkin_ws/src/uav/scripts/path_planner")
sys.path.insert(0, "/home/ninetripod/catkin_ws/src/uav/scripts/control")
sys.path.insert(0, "/home/ninetripod/catkin_ws/src/uav/scripts/sensor")

import rospy
import rplidar
from std_msgs.msg import String
import json, math, traceback
import numpy as np
import logging
from shape import Pillar
from path import Path
from filter import Filter

import serial
from link import Link
import matplotlib.pyplot as plt
from peripherals import get_LASER, mPWM
from sound import play_sound, play_target_audio
from fix import Projection
import time


class Navigater:
    TASK_NONE                   = 0
    TASK_WAIT_LEARNING_TARGET   = 1
    TASK_WAIT_TAKEOFF           = 2
    TASK_TAKING_OFF             = 3
    TASK_FIND_CIRCLE            = 4
    TASK_CALC_PATH              = 5
    TASK_SEARCHING_TARGET       = 6
    TASK_TRACKING_TARGET        = 7
    TASK_WAIT_DELIVERY          = 8
    TASK_DELIVERY               = 9
    TASK_FLY_THROUGH            = 10
    TASK_END                    = 11

    TASK_DEBUG                  = 99


    # LASER_PIXEL_X = 350 # inverted 
    # LASER_PIXEL_Y = 285

    # With distort
    LASER_PIXEL_X = 287 # inverted 
    LASER_PIXEL_Y = 220

    NORMAL_ALTITUDE = 128
    DELIVERY_ALTITUDE = 80
    CAMERA_FIX_COEF = 1.9 / 518



    def __init__(self) -> None:          
        rospy.init_node("uav_navigator")

        # debug
        # self._obstacle = [Pillar([120, 120], -45, 110)]
        # self._border = [-25, -25, 470, 300]
        # self._dense = 80
        # self._path_planner = Path(self._border, self._dense, self._obstacle)
        # self.paths = self._path_planner.Path_Global_Search(500)

        # self._path_planner.init_tracker()
        # self._path_planner.Path_Through_Ring([300.5, 300.2], [0, 0])
        # self._path_planner.draw()
        # print("path drawed")

        # init
        self._stick_finder = StickFinder()

        self._delay_count = 0
        self._task = Navigater.TASK_WAIT_LEARNING_TARGET

        self.link = Link(timeout = 0.05)
        
        self.laser = get_LASER() # ji guang

        self.cargo = Cargo()

        self.camera_fix = None

        # ros init : must be called the at end
        self.ladar_sub = rplidar.Rplidar_Subscriber()
        self.pos_sub = UavPosSubscriber()
        self.target_sub = TargetFinder()

        self.x_filter = Filter(2,3)
        self.y_filter = Filter(2,3)

        self.anchor_pos = [0,0]
        self.data_info = []

        self._timer = rospy.Timer(
            rospy.Duration(0.05),
            self.task
        )

        self.mode_pub = rospy.Publisher("uav/mode", String, queue_size = 5)
        play_sound("/home/ninetripod/catkin_ws/src/uav/resource/audio/begin_detect.mp3")
        # debug
        # play_sound("/home/ninetripod/catkin_ws/src/uav/resource/audio/sound_long.mp3", 20)
        # self.link.send(Link.pack_pos_global(100, 100, Navigater.NORMAL_ALTITUDE))

    # def set_task(self, value: int):
    #     self.__task = value
    #     self.mode_pub.publish(String(value))
    
    # def get_task(self):
    #     return self.__task

    def task(self, timer):

        if self._task == Navigater.TASK_WAIT_LEARNING_TARGET:
            if self.pos_sub.available and self._delay(10):
                self.mode_pub.publish(str(self._task))
                target = self.target_sub.target
                if target is not None:
                    play_target_audio(target)

                    rospy.loginfo(f"set target: {target}")
                    self.target_sub._data = None
                    self._task = Navigater.TASK_WAIT_TAKEOFF
                    self.mode_pub.publish(str(self._task))
        

        elif self._task == Navigater.TASK_WAIT_TAKEOFF:
            if self.pos_sub.available:
                pos = self.pos_sub.pos
                if pos[2] < 1:
                    self._task = Navigater.TASK_TAKING_OFF
                    self.mode_pub.publish(str(self._task))
                    rospy.loginfo(f"taking off")
                    play_sound("/home/ninetripod/catkin_ws/src/uav/resource/audio/fly_ready.mp3")
                    rospy.sleep(1)


        # wait to take off
        elif self._task == Navigater.TASK_TAKING_OFF:
            if self.pos_sub.available:
                pos = self.pos_sub.pos
                if self.camera_fix == None and pos[2] < 0.5:
                    self.camera_fix = Projection(self.pos_sub.orient)
                    print("init camera fix")
                    
                # print(pos)
                if pos[2] > Navigater.NORMAL_ALTITUDE - 10:
                    self.target_sub._data = None
                    self.target_sub.available = False

                    self._task = Navigater.TASK_FIND_CIRCLE
                    rospy.sleep(1)
                    rospy.loginfo("start finding circle")


        # find circle
        elif self._task == Navigater.TASK_FIND_CIRCLE:
            if self._delay(10):
                data = self.ladar_sub.ladar_data
                if data is not None:
                    data = data[0: 885]
                    rplidar.range_filter_rect(data, 350, 300)
                    res = self._stick_finder.find_circle(data)
                    if res is not None:
                        rospy.loginfo(f"find circle at {res}")
                        
                        p0 = res[0]
                        p1 = res[1]
                        self.circle_center = [(p0[0] + p1[0])/2, (p0[1] + p1[1])/2]
                        
                        # debug
                        # center = [(p0[0] + p1[0]) / 2, (p0[1] + p1[1]) / 2]
                        # msg = Link.pack_pos_global(center[0], center[1], Navigater.NORMAL_ALTITUDE)
                        # self.link.send(msg)
                        # self.ladar_sub.stop_motor()
                        # self._task = Navigater.TASK_NONE
                        # return
                    

                        circle = Pillar([120, 120], -45, 110)
                        circle.set_endpoint(p0, p1)
                        self._obstacle = [circle]
                        self._border = [-25, -25, 450, 300]
                        self._dense = 80
                        self._path_planner = Path(self._border, self._dense, self._obstacle)
                        self.paths = self._path_planner.Path_Global_Search()
                        
                        self._path_planner.draw("search",'/home/ninetripod/catkin_ws/src/uav/log/', save = True)
                        print("path drawed")
                        self.ladar_sub.stop_motor()

                        self.target_sub._data = None
                        self.target_sub.available = False
                        self._task = Navigater.TASK_SEARCHING_TARGET


        # searching
        elif self._task == Navigater.TASK_SEARCHING_TARGET:
            path = self._path_planner.Tracker(self.pos_sub.pos[:-1])
            msg = Link.pack_pos_global(path[0], path[1], extra = Navigater.NORMAL_ALTITUDE)
            self.link.send(msg)
            
            target = self.target_sub.target_pos

            if target is not None and self.target_sub.available:
                self.laser.write(True)
                self._task = Navigater.TASK_TRACKING_TARGET

                rospy.loginfo("start tracking target")

        
        elif self._task == Navigater.TASK_TRACKING_TARGET:
            pixel = self.target_sub.target_pos
            pos = self.pos_sub.pos

            # miss target for a long time or false target
            if time.time() - self.target_sub.last_find_time > 4:
                self._task = Navigater.TASK_SEARCHING_TARGET
                self.target_sub.available = False
                print("missing target for a long time")
                return

            if pixel is not None and pos is not None:
                pixel_offset_x = pixel[0] - Navigater.LASER_PIXEL_X # err bewteen laser and center
                pixel_offset_y = pixel[1] - Navigater.LASER_PIXEL_Y

                # camera coordinate to world coordinate
                world_offset_x = - pixel_offset_y * pos[2] * Navigater.CAMERA_FIX_COEF
                world_offset_y = - pixel_offset_x * pos[2] * Navigater.CAMERA_FIX_COEF

                # projection transform under world coordinate
                self.camera_fix.pre_transform(self.pos_sub.orient, pos[2])
                world_offset_pos = self.camera_fix.transform([world_offset_x, world_offset_y])

                filted_offset_x = self.x_filter.mean_filter_thres(world_offset_pos[0])
                filted_offset_y = self.y_filter.mean_filter_thres(world_offset_pos[1]) 
                target_x = filted_offset_x + pos[0]  # inverted, scale
                target_y = filted_offset_y + pos[1]

                if target_x < self._border[0]:
                    target_x = self._border[0]

                elif target_x > self._border[2]:
                    target_x = self._border[2]

                if target_y < self._border[1]:
                    target_y = self._border[1]

                elif target_y > self._border[3]:
                    target_y = self._border[3]

                # self.anchor_pos = [pos[0], pos[1]]


                # print(pixel, pos)

                # 飞控收到后会开始下降放货
                target = self._path_planner.Pos_Anti_Collision([pos[0], pos[1]], [target_x, target_y])

                # if self._delay(10):
                # print("tar:", target)

                self.data_info.append([target, pos])

                msg = Link.pack_pos_global(target[0], target[1], extra = Navigater.NORMAL_ALTITUDE)
                self.link.send(msg)

                # ===============================================================================================

                # self.anchor_pos = [target_x, target_y]
                # rospy.loginfo("anchor pos{}".format(self.anchor_pos))

                # self._path_planner.init_tracker()
                # self._path_planner.Path_Through_Ring([target_x, target_y], [0.0, 0.0])
                # self._path_planner.draw('through','/home/ninetripod/catkin_ws/src/uav/log/', save = True)

                # try:
                #     with open('/home/ninetripod/catkin_ws/src/uav/log/data.json','wt') as f:
                #         json.dump(self.data_info, f)
                # except:
                #     pass

                # rospy.sleep(5)
                # # play_sound()
                # self._task = Navigater.TASK_DELIVERY

                # ==================================================================================================

                offset_d = math.sqrt(filted_offset_x * filted_offset_x + filted_offset_y * filted_offset_y)

                if offset_d < 10.0:
                    self._task = Navigater.TASK_WAIT_DELIVERY
                    self.anchor_pos = [pos[0], pos[1]]
                    self._path_planner.init_tracker()
                    self._path_planner.Path_Through_Ring([pos[0], pos[1]], [0.0, 0.0])
                    self._path_planner.draw('through','/home/ninetripod/catkin_ws/src/uav/log/', save = True)

                    try:
                        with open('/home/ninetripod/catkin_ws/src/uav/log/data.json','wt') as f:
                            json.dump(self.data_info, f, indent = 4)
                    except:
                        pass

                    msg = Link.pack_pos_global(pos[0], pos[1], extra = Navigater.DELIVERY_ALTITUDE)
                    self.link.send(msg)

                    
                    # self.laser.write(False)

                    rospy.loginfo("start fly through obstacle")


        elif self._task == Navigater.TASK_WAIT_DELIVERY:
            pos = self.pos_sub.pos
            if pos[2] < Navigater.DELIVERY_ALTITUDE + 10:
                self._task = Navigater.TASK_DELIVERY
                rospy.sleep(1.5)
                
        
        elif self._task == Navigater.TASK_DELIVERY:
            if self.cargo.delivery():   # complete delivering
                # resume height
                # play_sound(amp = 10)
                msg = Link.pack_pos_global(self.anchor_pos[0], self.anchor_pos[1], Navigater.NORMAL_ALTITUDE)
                self.link.send(msg)
        
                self.laser.write(False)

                self._task = Navigater.TASK_FLY_THROUGH
    

        elif self._task == Navigater.TASK_FLY_THROUGH:
            pos = self.pos_sub.pos
            if pos[2] > Navigater.NORMAL_ALTITUDE - 10:
                path = self._path_planner.Tracker(pos[:-1])
                msg = Link.pack_pos_global(path[0], path[1], extra = Navigater.NORMAL_ALTITUDE)
                self.link.send(msg)

                d = math.sqrt(pos[0]*pos[0] + pos[1]*pos[1])
                if d < 3 or (path[0] == 0.0 and path[1] == 0.0):
                    self._task = Navigater.TASK_END


        elif self._task == Navigater.TASK_END:
            msg = Link.pack_pos_global(0, 0, extra = 0)
            self.link.send(msg)
            self._task = Navigater.TASK_NONE


        elif self._task == Navigater.TASK_DEBUG:
            # print(self.pos_sub.pos)
            # print(self.pos_sub.orient)
            # print()
            pass
                

    def _delay(self, tick):
        if self._delay_count < tick:
            self._delay_count += 1
            return False
        else:
            self._delay_count = 0
            return True

    def run(self):
        rospy.spin()


class TargetFinder():
    def __init__(self) -> None:
        rospy.Subscriber("uav/pixel_pos", String, self.onData)
        self._data = None
        self.last_find_time = time.time()
        self.available = False

    def onData(self, data):
        self._data = data.data
        self.last_find_time = time.time()
        self.available = True
        # print("tr updt", self._data)
    
    @property
    def target_pos(self):
        if self._data is None:
            return None
        else:
            return json.loads(self._data)

    @property
    def target(self):
        if self._data is None:
            return None
        else:
            return self._data
        
        
    


class UavPosSubscriber:
    def __init__(self) -> None:
        rospy.Subscriber("uav/pos", String, self.onData)
        self._data = None
        self.available = False

    def onData(self, data):
        self._data = data.data
        self.available = True
    
    @property
    def pos(self):
        if self._data is None:
            return [0, 0, 0]
        return json.loads(self._data)["p"]
    
    @property
    def orient(self):
        if self._data is None:
            return [0, 0, 0, 0]
        
        return json.loads(self._data)["o"]


class StickFinder():
    def __init__(self) -> None:
        # self._circle_pos = np.array([[0., 0.], [0., 0.]])
        self._find_obstacle_count = 0
        self.candidate = [] # [res, belief]
    
    def find_circle(self, data):
        c = rplidar.find_circle4(data, log = False)
        if c is not None:
            print(c)
            self._find_obstacle_count += 1

            if len(self.candidate) == 0:
                self.candidate.append([c, -1])

            believable = False
            for cand in self.candidate:
                var = np.var(cand[0] - c)
                if var < 8:
                    cand[0] = (cand[0] + c) / 2
                    cand[1] += 1
                    believable = True
            
            if not believable:
                self.candidate.append([c, 0])
        

        if self._find_obstacle_count == 8:
            self.candidate.sort(key = lambda c : c[1], reverse = True)
            print("find circle", self.candidate)
            return self.candidate[0][0]
        
        return None

class Cargo():
    # ADJUST_HEIGHT = 0
    SEND_DOWN_ORDER = 0
    WAIT_DOWN = 1
    HOVERING = 2
    SEND_UP_ORDER = 3
    WAIT_UP = 4

    def __init__(self) -> None:
        self.p = mPWM(0, 0)
        self._delay_count = 0
        self._cargo_status = Cargo.SEND_DOWN_ORDER
        print("init pwm", self.p)

    def up(self):
        self.p.period = 20_000_000         # 20ms
        self.p.duty_cycle = 1_000_000   # 1000us
        self.p.polarity = mPWM.POLARITY_NORMAL
        self.p.enable = 1

    def stop(self):
        self.p.enable = 0

    def down(self):
        self.p.period = 20_000_000         # 20ms
        self.p.duty_cycle = 2_000_000   # 2000us
        self.p.polarity = mPWM.POLARITY_NORMAL
        self.p.enable = 1

    def delivery(self) -> bool:
        if self._cargo_status == Cargo.SEND_DOWN_ORDER:
            play_sound(amp = 30)
            rospy.loginfo("start play audio")
            self.down()
            self._cargo_status = Cargo.WAIT_DOWN
            rospy.loginfo("start delivering")

        elif self._cargo_status == Cargo.WAIT_DOWN:
            if self._delay(50):
                self.stop()
                self._cargo_status = Cargo.HOVERING
                
        elif self._cargo_status == Cargo.HOVERING:
            if self._delay(80):
                self._cargo_status = Cargo.SEND_UP_ORDER
        
        elif self._cargo_status == Cargo.SEND_UP_ORDER:
            self.up()
            self._cargo_status =  Cargo.WAIT_UP
        
        elif self._cargo_status == Cargo.WAIT_UP:
            if self._delay(70):
                self.stop()
                self.p.close()
                rospy.loginfo("stop delivering")
                return True
        
        return False

    def _delay(self, tick):
        if self._delay_count < tick:
            self._delay_count += 1
            return False
        else:
            self._delay_count = 0
            return True


if __name__ == "__main__":
    # logging.basicConfig(filename="/home/ninetripod/catkin_ws/src/uav/log.log", level = logging.NOTSET)
    try:
        nav = Navigater()
        nav.run()
    
    except Exception as e:
        traceback.print_exc()
    
    finally:
        nav.laser.write(False)
        nav.laser.close()
        nav.cargo.p.close()